package brave.spring.rabbit;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.internal.Nullable;
import brave.messaging.MessagingRequest;
import brave.propagation.TraceContext;
import brave.propagation.TraceContextOrSamplingFlags;
import brave.sampler.SamplerFunction;
import com.rabbitmq.client.AMQP;
import com.rabbitmq.client.Envelope;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessagePostProcessor;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.amqp.rabbit.core.RabbitTemplate;

import static brave.Span.Kind.CONSUMER;
import static brave.spring.rabbit.SpringRabbitTracing.*;

/**
 * 接收后跟踪处理器
 *
 * @author zsy
 * @see brave.spring.rabbit.SpringRabbitTracing
 * @see brave.spring.rabbit.TracingMessagePostProcessor
 * @see brave.spring.rabbit.TracingRabbitListenerAdvice
 * @since 2023/8/15
 */
public class TracingRabbitPostProcessor implements MessagePostProcessor {

    final ThreadLocal<SpanScope> local;
    final SpringRabbitTracing springRabbitTracing;
    final Tracing tracing;
    final Tracer tracer;
    final TraceContext.Extractor<MessageConsumerRequest> extractor;
    final SamplerFunction<MessagingRequest> sampler;
    @Nullable
    final String remoteServiceName;

    public TracingRabbitPostProcessor(SpringRabbitTracing springRabbitTracing) {
        this.springRabbitTracing = springRabbitTracing;
        this.tracing = springRabbitTracing.tracing;
        this.tracer = tracing.tracer();
        this.extractor = springRabbitTracing.consumerExtractor;
        this.sampler = springRabbitTracing.consumerSampler;
        this.remoteServiceName = springRabbitTracing.remoteServiceName;
        this.local = new ThreadLocal<>();
    }

    @Override
    public Message postProcessMessage(Message message) throws AmqpException {

        MessageConsumerRequest request = new MessageConsumerRequest(message);

        TraceContextOrSamplingFlags extracted =
                springRabbitTracing.extractAndClearTraceIdHeaders(extractor, request, message);

        // named for BlockingQueueConsumer.nextMessage, which we can't currently see
        Span consumerSpan = springRabbitTracing.nextMessagingSpan(sampler, request, extracted);
        Span listenerSpan = tracer.newChild(consumerSpan.context());

        if (!consumerSpan.isNoop()) {
            setConsumerSpan(consumerSpan, message.getMessageProperties());

            // incur timestamp overhead only once
            long timestamp = tracing.clock(consumerSpan.context()).currentTimeMicroseconds();
            consumerSpan.start(timestamp);
            long consumerFinish = timestamp + 1L; // save a clock reading
            consumerSpan.finish(consumerFinish);

            // not using scoped span as we want to start with a pre-configured time
            listenerSpan.name("on-message").start(consumerFinish);
        }
        /**
         * 在使用 rabbitTemplate.receive 方法，链路会断开，封装的方法无法做到对应操作，可以考虑自行扩展手写
         * 找到绑定的切入点，一般为aop最佳，否则无法做到关闭scope， 只能勉强做到绑定，但是没有最大close地点
         * {@link RabbitTemplate#doReceiveNoWait(String)}
         * {@link RabbitTemplate#buildMessage(Envelope, AMQP.BasicProperties, byte[], int)}
         */
        set(listenerSpan);
        return message;
    }

    /**
     * 参考rabbit brave释放过程
     * Throwable error = null;
     * try {
     * return message;
     * } catch (Throwable t) {
     * error = t;
     * throw t;
     * } finally {
     * if (error != null) listenerSpan.error(error);
     * listenerSpan.finish();
     * ws.close();
     * }
     *
     * @param spanScope
     * @param error
     */
    protected void release(SpanScope spanScope, Throwable error) {
        if (spanScope != null) {
            if (error != null) spanScope.span.error(error);
            spanScope.span.finish();
            spanScope.scope.close();
        }
    }

    protected void set(Span listenerSpan) {
        release(local.get(), null);
        local.set(new SpanScope(listenerSpan, tracer.withSpanInScope(listenerSpan)));
    }

    public void close(Throwable error) {
        release(local.get(), error);
        local.remove();
    }

    void setConsumerSpan(Span span, MessageProperties properties) {
        span.name("next-message").kind(CONSUMER);
        maybeTag(span, RABBIT_EXCHANGE, properties.getReceivedExchange());
        maybeTag(span, RABBIT_ROUTING_KEY, properties.getReceivedRoutingKey());
        maybeTag(span, RABBIT_QUEUE, properties.getConsumerQueue());
        if (remoteServiceName != null) span.remoteServiceName(remoteServiceName);
    }

    static void maybeTag(Span span, String tag, String value) {
        if (value != null) span.tag(tag, value);
    }

    public record SpanScope(Span span, Tracer.SpanInScope scope) {
    }

}
